import os
from pinecone import Pinecone, ServerlessSpec
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

# 初始化Pinecone
def setup_pinecone():
    # 使用你的API密钥
    api_key = 'pcsk_3feH1Q_FWFLXijCNDKBZXWUyH2KXTNRCV67CBXCc6f4DjqmtBaD2XrBjdouZ18FPDxjqJA'
    
    # 初始化Pinecone
    pc = Pinecone(api_key=api_key)
    
    # 列出所有索引
    indexes = pc.list_indexes()
    print("现有索引:", [index.name for index in indexes.indexes])
    
    return pc

def create_example_index():
    """创建示例索引并演示基本操作"""
    pc = setup_pinecone()
    
    # 索引名称
    index_name = "mnist-example"
    
    # 如果索引已存在，先删除
    existing_indexes = pc.list_indexes()
    if index_name in [index.name for index in existing_indexes.indexes]:
        print(f"删除已存在的索引: {index_name}")
        pc.delete_index(index_name)
    
    # 创建新索引
    print(f"创建新索引: {index_name}")
    pc.create_index(
        name=index_name,
        dimension=8,  # 为了示例使用较小的维度
        metric="cosine",
        spec=ServerlessSpec(
            cloud='aws',
            region='us-east-1'
        )
    )
    
    # 连接到索引
    index = pc.Index(index_name)
    
    # 生成一些示例数据
    print("\n生成示例数据...")
    vectors = []
    for i in range(10):
        # 创建随机向量
        vector = np.random.randn(8).tolist()
        vectors.append({
            "id": f"vec_{i}", 
            "values": vector, 
            "metadata": {"label": f"digit_{i % 3}"}
        })
    
    # 上传向量
    print("上传向量到Pinecone...")
    index.upsert(vectors=vectors)
    
    # 查询示例
    print("\n执行查询...")
    query_vector = np.random.randn(8).tolist()
    results = index.query(
        vector=query_vector,
        top_k=3,
        include_metadata=True
    )
    
    print("查询结果:")
    for match in results['matches']:
        print(f"ID: {match['id']}, Score: {match['score']:.4f}, Label: {match['metadata']['label']}")
    
    # 清理：删除索引
    pc.delete_index(index_name)
    print(f"\n已清理索引: {index_name}")

if __name__ == "__main__":
    print("=== Pinecone 示例演示 ===")
    create_example_index()
    print("\n示例完成！")